package com.bins.langchain.easy.chroma.service.impl;

import com.bins.langchain.easy.chroma.service.ChromaService;
import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.vectorstore.ChromaVectorStore;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
public class ChromaServiceImpl implements ChromaService {

    @Resource
    private OllamaEmbeddingModel ollamaEmbeddingModel;

    @Resource
    private ChromaVectorStore chromaVectorStore;

    @Override
    public void saveEmbedding() {
        List<Document> documents = List.of(
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
                new Document("The World is Big and Salvation Lurks Around the Corner"),
                new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));
        chromaVectorStore.add(documents);
    }

    @Override
    public List<String> search() {
        List<Document> results = chromaVectorStore.similaritySearch("Spring");
        if (CollectionUtils.isEmpty(results)) {
            return Arrays.asList("无数据");
        }
        return results.stream().map(Document::getContent).collect(Collectors.toList());
    }
}
